//! unmp子协议

#![no_std]
extern crate alloc;

use alloc::boxed::Box;
use alloc::collections::BTreeMap;
use async_ach_mpmc::heap as mpmc;
use futures_util::future::{select, select_all, Either};
pub use unmp_net::packet::{Bytes, Id, Packet};
pub use unmp_net::Network;

pub struct Data {
    pub remote: Id,
    pub data: Bytes,
}

pub struct DataChannel {
    pub tx: mpmc::Sender<Data, 8, 3, 1>,
    pub rx: mpmc::Receiver<Data, 8, 3, 1>,
}
impl DataChannel {
    pub async fn send(&self, packet: Data) {
        self.tx.send(packet).await
    }
    pub async fn recv(&self) -> Data {
        self.rx.recv().await
    }
    pub fn split(self) -> (mpmc::Sender<Data, 8, 3, 1>, mpmc::Receiver<Data, 8, 3, 1>) {
        let DataChannel { tx, rx } = self;
        (tx, rx)
    }
}

pub struct Distributor {
    net: Network,
    chs: BTreeMap<u8, DataChannel>,
}
impl Distributor {
    pub fn new(net: Network) -> Self {
        Self {
            net,
            chs: BTreeMap::new(),
        }
    }
    /// return None if the protocol had registered
    pub fn register(&mut self, protocol: u8) -> Option<DataChannel> {
        if self.chs.contains_key(&protocol) {
            return None;
        }
        let (tx_tx, tx_rx) = mpmc::channel();
        let (rx_tx, rx_rx) = mpmc::channel();
        self.chs.insert(
            protocol,
            DataChannel {
                tx: rx_tx,
                rx: tx_rx,
            },
        );
        Some(DataChannel {
            tx: tx_tx,
            rx: rx_rx,
        })
    }
    pub async fn daemon(mut self) -> ! {
        loop {
            let recv_net = Box::pin(self.net.recv());
            let recv_tx = select_all(self.chs.iter_mut().map(|(_, ch)| Box::pin(ch.recv())));
            let result = match select(recv_net, recv_tx).await {
                Either::Left((packet, _)) => Either::Left(packet),
                Either::Right(((data, i, _), _)) => Either::Right((data, i)),
            };
            match result {
                Either::Left(packet) => {
                    if let Some(ch) = self.chs.get_mut(&packet.head.protocol) {
                        let data = Data {
                            remote: packet.head.src,
                            data: packet.data,
                        };
                        ch.send(data).await;
                    }
                }
                Either::Right((data, i)) => {
                    if let Some(protocol) = self.chs.keys().nth(i) {
                        let packet = self.net.new_packet(data.remote, *protocol, data.data);
                        let _ = self.net.send(packet).await;
                    }
                }
            };
        }
    }
}
impl Drop for Distributor {
    fn drop(&mut self) {
        panic!("please call daemon of Distributor");
    }
}
